#include "light_corner_correct.h"
#include <numeric>

LightCornerCorrector::LightCornerCorrector() {
    // 构造函数实现可以为空
}

void LightCornerCorrector::correctCorners(Armor &armor, const cv::Mat &gray_img) {
    // If the width of the light is too small, the correction is not performed
    constexpr int PASS_OPTIMIZE_WIDTH = 3;

    if (armor.left_light.width > PASS_OPTIMIZE_WIDTH) {
        // Find the symmetry axis of the light
        SymmetryAxis left_axis = findSymmetryAxis(gray_img, armor.left_light);
        armor.left_light.center = left_axis.centroid;
        armor.left_light.axis = left_axis.direction;
        // Find the corner of the light
        cv::Point2f t = findCorner(gray_img, armor.left_light, left_axis, "top");
        if (t.x > 0) {
            armor.left_light.top = t;
        }
        cv::Point2f b = findCorner(gray_img, armor.left_light, left_axis, "bottom");
        if (b.x > 0) {
            armor.left_light.bottom = b;
        }
    }

    if (armor.right_light.width > PASS_OPTIMIZE_WIDTH) {
        // Find the symmetry axis of the light
        SymmetryAxis right_axis = findSymmetryAxis(gray_img, armor.right_light);
        armor.right_light.center = right_axis.centroid;
        armor.right_light.axis = right_axis.direction;
        // Find the corner of the light
        cv::Point2f t = findCorner(gray_img, armor.right_light, right_axis, "top");
        if (t.x > 0) {
            armor.right_light.top = t;
        }
        cv::Point2f b = findCorner(gray_img, armor.right_light, right_axis, "bottom");
        if (b.x > 0) {
            armor.right_light.bottom = b;
        }
    }
}

SymmetryAxis LightCornerCorrector::findSymmetryAxis(const cv::Mat &gray_img, const LightBar &light) {
    constexpr float MAX_BRIGHTNESS = 25;
    constexpr float SCALE = 0.07;

    // Scale the bounding box
    cv::Rect light_box = light.boundingRect();
    light_box.x -= light_box.width * SCALE;
    light_box.y -= light_box.height * SCALE;
    light_box.width += light_box.width * SCALE * 2;
    light_box.height += light_box.height * SCALE * 2;

    // Check boundary
    light_box.x = std::max(light_box.x, 0);
    light_box.x = std::min(light_box.x, gray_img.cols - 1);
    light_box.y = std::max(light_box.y, 0);
    light_box.y = std::min(light_box.y, gray_img.rows - 1);
    light_box.width = std::min(light_box.width, gray_img.cols - light_box.x);
    light_box.height = std::min(light_box.height, gray_img.rows - light_box.y);

    // Get normalized light image
    cv::Mat roi = gray_img(light_box);
    float mean_val = (float)cv::mean(roi)[0];

    // 优化：直接使用整数运算避免浮点运算
    cv::Mat roi_normalized;
    roi.convertTo(roi_normalized, CV_8U, 255.0 / MAX_BRIGHTNESS / 255.0);

    // Calculate the centroid
    cv::Moments moments = cv::moments(roi_normalized, false);
    cv::Point2f centroid = cv::Point2f((float)(moments.m10 / moments.m00), (float)(moments.m01 / moments.m00)) +
                           cv::Point2f((float)light_box.x, (float)light_box.y);

    // 优化：减少点云生成的数量，使用采样而不是每个像素都添加点
    std::vector<cv::Point2f> points;
    // 降低采样率以提高性能
    const int sample_rate = 2; // 每隔2个像素采样一个点
    for (int i = 0; i < roi_normalized.rows; i += sample_rate) {
        for (int j = 0; j < roi_normalized.cols; j += sample_rate) {
            uchar value = roi_normalized.at<uchar>(i, j);
            // 根据亮度值添加相应数量的点，但限制最大数量
            int count = std::min(value / 5, 5); // 限制每个像素最多添加5个点
            for (int k = 0; k < count; k++) {
                points.emplace_back(cv::Point2f((float)j, (float)i));
            }
        }
    }

    if (points.empty()) {
        // 如果没有点，使用默认方向
        cv::Point2f axis(0, -1); // 默认向上
        return SymmetryAxis{centroid, axis, mean_val};
    }

    cv::Mat points_mat = cv::Mat(points).reshape(1);

    // PCA (Principal Component Analysis)
    cv::PCA pca(points_mat, cv::Mat(), cv::PCA::DATA_AS_ROW);

    // Get the symmetry axis
    cv::Point2f axis = cv::Point2f(pca.eigenvectors.at<float>(0, 0), pca.eigenvectors.at<float>(0, 1));

    // Normalize the axis
    float norm = std::sqrt(axis.x * axis.x + axis.y * axis.y);
    if (norm > 0) {
        axis.x /= norm;
        axis.y /= norm;
    }

    if (axis.y > 0) {
        axis.x = -axis.x;
        axis.y = -axis.y;
    }

    return SymmetryAxis{centroid, axis, mean_val};
}

cv::Point2f LightCornerCorrector::findCorner(const cv::Mat &gray_img,
                                             const LightBar &light,
                                             const SymmetryAxis &axis,
                                             const std::string &order) {
    constexpr float START = 0.8f / 2;
    constexpr float END = 1.2f / 2;

    auto inImage = [&gray_img](const cv::Point &point) -> bool {
        return point.x >= 0 && point.x < gray_img.cols && point.y >= 0 && point.y < gray_img.rows;
    };

    auto distance = [](float x0, float y0, float x1, float y1) -> float {
        // 优化：使用近似距离计算，避免开方运算
        float dx = x0 - x1;
        float dy = y0 - y1;
        return std::abs(dx) + std::abs(dy); // 曼哈顿距离作为近似
    };

    int oper = order == "top" ? 1 : -1;
    float L = (float)light.length;
    float dx = axis.direction.x * oper;
    float dy = axis.direction.y * oper;

    std::vector<cv::Point2f> candidates;

    // 优化：减少候选点数量
    int n = (int)light.width - 2;
    int half_n = std::round(n / 2.0f);

    // 限制搜索点的数量
    const int max_search_points = 3;
    int step = std::max(1, half_n / max_search_points);

    for (int i = -half_n; i <= half_n; i += step) {
        float x0 = axis.centroid.x + L * START * dx + i;
        float y0 = axis.centroid.y + L * START * dy;

        cv::Point2f prev = cv::Point2f(x0, y0);
        cv::Point2f corner = cv::Point2f(x0, y0);
        float max_brightness_diff = 0;
        bool has_corner = false;

        // 优化：限制搜索步数
        const int max_steps = 20;
        int steps = 0;

        // Search along the symmetry axis to find the corner that has the maximum brightness difference
        for (float x = x0 + dx, y = y0 + dy;
             distance(x, y, x0, y0) < L * (END - START) && steps < max_steps;
             x += dx, y += dy, steps++) {
            cv::Point2f cur = cv::Point2f(x, y);
            if (!inImage(cv::Point(cur))) {
                break;
            }

            // 优化：使用边界检查避免访问无效内存
            if (prev.y >= 0 && prev.y < gray_img.rows && prev.x >= 0 && prev.x < gray_img.cols &&
                cur.y >= 0 && cur.y < gray_img.rows && cur.x >= 0 && cur.x < gray_img.cols) {

                float brightness_diff = (float)(gray_img.at<uchar>((int)prev.y, (int)prev.x) -
                                              gray_img.at<uchar>((int)cur.y, (int)cur.x));
                if (brightness_diff > max_brightness_diff && gray_img.at<uchar>((int)prev.y, (int)prev.x) > axis.mean_val) {
                    max_brightness_diff = brightness_diff;
                    corner = prev;
                    has_corner = true;
                }
            }

            prev = cur;
        }

        if (has_corner) {
            candidates.emplace_back(corner);
        }
    }

    if (!candidates.empty()) {
        cv::Point2f result = std::accumulate(candidates.begin(), candidates.end(), cv::Point2f(0, 0));
        return result / (float)candidates.size();
    }

    return cv::Point2f(-1, -1);
}
